# -*- coding: UTF-8 -*-
"""
@项目名称：kimi_get_protein.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-19-09:14
"""
from Bio.PDB import PDBParser, PDBIO, is_aa  # 导入is_aa函数用于判断标准氨基酸残基[2](@ref)
import os


def extract_and_save_chains(pdb_file_path, output_dir, only_aa=True):
    """
    读取PDB文件，并将其中的每条链单独保存为新的PDB文件。
    :param pdb_file_path: 输入的PDB文件路径
    :param output_dir: 输出目录，用于保存各链的PDB文件
    """
    # 创建PDB解析器，QUIET=True用于忽略警告信息
    parser = PDBParser(QUIET=True)
    # 解析PDB文件，获取结构对象
    # 'structure_id'可以是任意标识符，通常用PDB代码
    structure = parser.get_structure('protein', pdb_file_path)  #
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    # 创建PDBIO对象用于保存结构
    io = PDBIO()
    for model in structure:
        for chain in model:
            chain_id = chain.get_id()

            # 生成输出文件名
            output_filename = f"{chain_id}.pdb"
            output_path = os.path.join(output_dir, output_filename)
            if only_aa and not any(is_aa(residue) for residue in chain): # 检查链中是否至少有一个标准氨基酸残基
                # 当仅保留蛋白质链时，不符合要求的跳过
                # 如果你只想保存蛋白质链排除水分子、离子等），可以在保存前检查链中是否包含氨基酸残基。
                chain_id = chain.get_id()
                print(f"链 {chain_id} 不包含标准氨基酸残基，已跳过。")
                continue
            io.set_structure(chain)
            io.save(output_path)
            print(f"已保存蛋白质链 {chain_id} 到文件 {output_path}")

if __name__ == "__main__":
    pdb_file = "/4T/AIDDserver/DataSrc/lpa_test/4bvw/4bvw.pdb"  # 请替换为你的PDB文件路径
    output_directory = f"{os.path.dirname(pdb_file)}/extracted_chains_all"  # 输出目录名称

    extract_and_save_chains(pdb_file, output_directory,only_aa=False)
